package com.zjl.media.core.javacv;

import com.zjl.media.config.properties.MediaProperties;
import com.zjl.media.core.MediaTransferContext;
import com.zjl.media.core.MediaTransferTask;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.ffmpeg.global.avcodec;
import org.bytedeco.ffmpeg.global.avutil;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.FFmpegFrameRecorder;
import org.bytedeco.javacv.FFmpegLogCallback;
import org.bytedeco.javacv.Frame;
import org.springframework.util.CollectionUtils;

import java.io.ByteArrayOutputStream;

/**
 * @author linfuxin Created on 2022-03-01 16:35:29
 */
@Slf4j
@Getter
public class JavacvTask implements MediaTransferTask {

    private final JavacvContext context;
    private final MediaProperties mediaProperties;

    public JavacvTask(JavacvContext mediaTransferContext, MediaProperties mediaProperties) {
        this.context = mediaTransferContext;
        this.mediaProperties = mediaProperties;
    }


    @Override
    public void run() {
        avutil.av_log_set_level(avutil.AV_LOG_ERROR);

        FFmpegLogCallback.set();
        try (
                FFmpegFrameGrabber grabber = createGrabber(context, mediaProperties);
                ByteArrayOutputStream output = context.getOutput();
        ) {
            grabber.start();
            log.info("拉流器启动成功!");

            try (FFmpegFrameRecorder recorder = createRecorder(output, grabber)) {
                recorder.start();
                log.info("推流器启动成功!");
                grabber.flush();
                if (context.getHeader() == null) {
                    context.setHeader(output.toByteArray());
                    context.getHeaderLatch().countDown();
                    for (ChannelHandlerContext client : context.getClients()) {
                        client.writeAndFlush(new BinaryWebSocketFrame(Unpooled.copiedBuffer(context.getHeader())));
                    }
                    output.reset();
                }
                recordAndSendFrame(grabber, recorder);
            }
        } catch (Exception e) {
            log.error("推流出现异常!", e);
        } finally {
            closeClient();
        }
    }

    private void recordAndSendFrame(FFmpegFrameGrabber grabber, FFmpegFrameRecorder recorder) {
        // 时间戳计算
        long startTime = 0;
        long videoTs;
        int maxRetryTimes = 5;
        ByteArrayOutputStream output = context.getOutput();
        for (int noFrameCount = 0, errorCount = 0; noFrameCount < maxRetryTimes && errorCount < maxRetryTimes; ) {
            try {
                // 转码
                long startGrabTime = System.currentTimeMillis();
                Frame frame = grabber.grab();
                if ((System.currentTimeMillis() - startGrabTime) > 3000) {
                    log.error("拉流器视频流时间戳异常!");
                    break;
                }
                if (frame == null) {
                    noFrameCount++;
                    continue;
                }
                if (startTime == 0) {
                    startTime = System.currentTimeMillis();
                }
                videoTs = 1000 * (System.currentTimeMillis() - startTime);
                if (videoTs > recorder.getTimestamp()) {
                    recorder.setTimestamp((videoTs));
                }
                recorder.record(frame);
                // 将缓存空间的引用计数-1，并将Packet中的其他字段设为初始值。如果引用计数为0，自动的释放缓存空间。
//                av_packet_unref(pkt);
            } catch (FFmpegFrameGrabber.Exception | FFmpegFrameRecorder.Exception e) {
                log.warn("转码异常!", e);
                errorCount++;
            }
            if (output.size() > 0) {
                byte[] bytes = output.toByteArray();
                output.reset();
                sendFrame(bytes);
            }
        }
        log.info("推流任务结束!");
    }

    private void sendFrame(byte[] data) {
        for (ChannelHandlerContext client : context.getClients()) {
            if (!client.channel().isWritable()) {
                // 直接跳过
                continue;
            }
            client.writeAndFlush(new BinaryWebSocketFrame(Unpooled.copiedBuffer(data)));
        }
    }

    private FFmpegFrameGrabber createGrabber(MediaTransferContext mediaTransferContext, MediaProperties mediaProperties) {
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(mediaTransferContext.getRtsp());
        grabber.setOptions(mediaProperties.getJavacv().getGrabberOption());
        return grabber;
    }

    private void closeClient() {
        if (context == null || CollectionUtils.isEmpty(context.getClients())) {
            return;
        }
        for (ChannelHandlerContext client : context.getClients()) {
            client.close();
        }
        context.setClients(null);
    }

    private FFmpegFrameRecorder createRecorder(ByteArrayOutputStream bos, FFmpegFrameGrabber grabber) {
        FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(bos, grabber.getImageWidth(), grabber.getImageHeight(),
                grabber.getAudioChannels());
        recorder.setFormat("flv");
        recorder.setInterleaved(false);
        recorder.setVideoOption("tune", "zerolatency");
        recorder.setVideoOption("preset", "ultrafast");
        recorder.setVideoOption("crf", "26");
        recorder.setVideoOption("threads", "1");
        // 视频帧率(保证视频质量的情况下最低25，低于25会出现闪屏)
        recorder.setFrameRate(25);
        // 关键帧间隔，一般与帧率相同或者是视频帧率的两倍
        recorder.setGopSize(25);
        recorder.setVideoBitrate(grabber.getVideoBitrate());

        recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
        recorder.setPixelFormat(avutil.AV_PIX_FMT_YUV420P);
        recorder.setAudioCodec(avcodec.AV_CODEC_ID_AAC);
        recorder.setTrellis(1);
        // 设置延迟
        recorder.setMaxDelay(0);

        return recorder;
    }
}
